"""
Asr模块，基于讯飞星火大模型的语音转文本的使用脚本。在本项目中主要的使用方法主要是将mp3的音频文件以流的形式发送到
语音转文本的服务器段，并不断返回转录的结果，
这里直接的转录结果会包含很多之前转录的重复内容（我使用了追加字符串的形式，所以有点问题）
问了解决这个转录结果字符串重复的问题，首先使用字符串操作先简单的清理优化字符串
核心是使用了大模型将文本进行优化得到最终的精准的转录文本结果
"""


import json
import hashlib
import websocket
import hmac
import base64
from urllib.parse import quote
import pydub
import numpy as np
import threading
import time
from collections import OrderedDict
import os
import traceback


class AudioTranscriber:
    def __init__(self):
        self.segments = OrderedDict()
        self.ws = None
        self.is_running = True
        self.last_update_time = time.time()

        self.API_KEY = "e1f21403934c092ac09416a353f5163e"
        self.APP_ID = "32863c5e"
        self.WS_URL = "ws://rtasr.xfyun.cn/v1/ws"

        #服务器转录基本参数
        self.SAMPLE_RATE = 16000    #采样率
        self.CHANNELS = 1
        self.BLOCK_SIZE = 2560
        self.SEND_INTERVAL = 0.025
        self.SILENCE_TIMEOUT = 3
        self.MIN_SEGMENT_LENGTH = 8
        self.MAX_OVERLAP_CHECK = 15

    def connect(self):
        #根据官方的标签生成方式，连接服务器是首先生成url标签
        ts = str(int(time.time()))
        base_str = (self.APP_ID + ts).encode('utf-8')
        md5 = hashlib.md5(base_str).hexdigest().encode('utf-8')
        signa = base64.b64encode(hmac.new(self.API_KEY.encode('utf-8'), md5, hashlib.sha1).digest()).decode('utf-8')
        url = f"{self.WS_URL}?appid={self.APP_ID}&ts={ts}&signa={quote(signa)}"

        self.ws = websocket.create_connection(url, timeout=10)
        threading.Thread(target=self.receiveMessage, daemon=True).start()

    def transcribe_mp3(self, file_path):
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")

            audio = pydub.AudioSegment.from_file(file_path)
            audio = audio.set_frame_rate(self.SAMPLE_RATE).set_channels(self.CHANNELS)
            samples = np.array(audio.get_array_of_samples()).astype(np.float32) / (2 ** 15)

            self.connect()
            time.sleep(1)

            for i in range(0, len(samples), self.BLOCK_SIZE):
                if not self.is_running: break
                self.sendAudio(samples[i:i + self.BLOCK_SIZE])
                time.sleep(self.SEND_INTERVAL)

            self.sendEndTag()
            #print("\n已发送结束标记，处理结果中...")

            wait_start = time.time()
            while time.time() - wait_start < self.SILENCE_TIMEOUT:
                if time.time() - self.last_update_time > 1: break
                time.sleep(0.1)

            final_text = self.getFinalText()
            return final_text

        except Exception as e:
            #print(f"ASR转录过程中发生严重错误: {e}")
            traceback.print_exc()
            return ""

        finally:
            self.close()

    #
    def sendAudio(self, audio_data):
        if self.ws and self.ws.connected:
            self.ws.send((audio_data * 32767).astype(np.int16).tobytes(), websocket.ABNF.OPCODE_BINARY)

    def sendEndTag(self):
        if self.ws and self.ws.connected:
            self.ws.send('{"end": true}')

    #处理转录文本内容的字符串操作，粗略优化转录文本
    def receiveMessage(self):
        try:
            while self.is_running and self.ws.connected:
                msg = self.ws.recv()
                if not msg: break
                result = json.loads(msg)

                if result.get("action") == "result":
                    self.processResult(result)
                elif result.get("action") == "error":
                    print(f"\nAPI错误: {result}")
                    self.close()
        except (websocket.WebSocketConnectionClosedException, ConnectionResetError):
            print("\n连接已关闭。")
        except Exception as e:
            print(f"\n接收消息时发生错误: {e}")
        finally:
            self.is_running = False

    def processResult(self, result):
        try:
            data = json.loads(result["data"])
            seg_id = data.get("seg_id")
            text = self.cleanText(data)

            if seg_id is not None and len(text) >= self.MIN_SEGMENT_LENGTH:
                if seg_id not in self.segments or not self.isRedundant(seg_id, text):
                    self.segments[seg_id] = text
                    self.last_update_time = time.time()
        except json.JSONDecodeError as e:
            print(f"\nJSON解析错误: {e}, 原始数据: {result['data']}")

    def isRedundant(self, seg_id, new_text):#判断冗余
        if seg_id not in self.segments: return False
        old_text = self.segments[seg_id]
        if new_text == old_text: return True
        if new_text in old_text and len(new_text) < len(old_text): return True
        if old_text in new_text and len(new_text) > len(old_text): return False
        overlap_length = self.findOverlapLen(old_text, new_text)
        min_length = min(len(old_text), len(new_text))
        return min_length > 0 and overlap_length / min_length > 0.7

    def findOverlapLen(self, text1, text2):
        max_possible = min(len(text1), len(text2), self.MAX_OVERLAP_CHECK)
        for i in range(max_possible, 0, -1):
            if text1.endswith(text2[:i]): return i
        return 0

    def cleanText(self, data):
        text = []
        try:
            for seg in data.get("cn", {}).get("st", {}).get("rt", []):
                for ws in seg.get("ws", []):
                    for cw in ws.get("cw", []):
                        word = cw.get("w", "")
                        if cw.get("wp", "") != "p":
                            text.append(word)
        except Exception:
            return ""
        return "".join(text)

    def getFinalText(self):
        self.mergeSegments()
        return self.transcription  # 暂时不进行后处理，以保留原始拼接结果

    def mergeSegments(self):
        if not self.segments:
            self.transcription = ""
            return

        sorted_segments = sorted(self.segments.items(), key=lambda x: x[0])
        merged_text = []
        last_segment = ""

        for seg_id, text in sorted_segments:
            if not text: continue
            if not last_segment:
                merged_text.append(text)
                last_segment = text
                continue

            overlap = self.findOverlapLen(last_segment, text)
            new_part = text[overlap:] if overlap > 0 else text

            if new_part:
                merged_text.append(new_part)
                last_segment = text

        self.transcription = "".join(merged_text)

    def close(self):
        self.is_running = False
        if self.ws and self.ws.connected:
            try:
                self.ws.close()
            except Exception:
                pass



if __name__ == "__main__":
    transcriber = AudioTranscriber()
    try:
        mp3_file = "test.mp3"
        result = transcriber.transcribe_mp3(mp3_file)
        print("=" * 60)
        print(result)
        print("=" * 60)
    except Exception as e:
        print(f"转录失败: {e}")
    finally:
        transcriber.close()
